# check_number.py
# python .\photo\check_number.py --db d:\baseline\photo\file_hashes.db --scan-dir D:\baseline\photo\

# 编写python程序 check_number.py --db sqlite数据库位置 --scan-dir 文件夹 
## --db sqlite数据库执行以下脚本: select zip_path,zip_name,count(*) as db_num from photo_hash group by zip_path,zip_name order by zip_path,zip_name
## --scan-dir 文件夹  下的zip_path目录下，找zip_name文件，得到这个zip文件下的文件数量 file_num
## 如果db_num 和file_num不一致，或者--scan-dir 文件夹不存在对应的zip文件，那么就屏幕显示出来

import argparse
import os
import sqlite3
import zipfile

def get_args():
    parser = argparse.ArgumentParser(description="检查数据库中记录的 ZIP 文件条目数量与实际 ZIP 文件中文件数量是否一致")
    parser.add_argument('--db', required=True, help='SQLite 数据库文件路径')
    parser.add_argument('--scan-dir', required=True, help='ZIP 文件所在的根目录')
    return parser.parse_args()

def get_zip_info_from_db(db_path):
    """从数据库中获取 zip_path, zip_name 和记录数量"""
    if not os.path.exists(db_path):
        raise FileNotFoundError(f"数据库文件不存在: {db_path}")

    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    query = """
        SELECT zip_path, zip_name, COUNT(*) AS db_num
        FROM photo_hash
        GROUP BY zip_path, zip_name
        ORDER BY zip_path, zip_name
    """

    try:
        cursor.execute(query)
        rows = cursor.fetchall()
        conn.close()
        return [(row[0], row[1], row[2]) for row in rows]  # (zip_path, zip_name, db_num)
    except Exception as e:
        conn.close()
        raise RuntimeError(f"执行数据库查询失败: {e}")

def count_zip_files(zip_full_path):
    """统计 ZIP 文件中包含的文件数量"""
    if not os.path.exists(zip_full_path):
        return None  # 文件不存在

    try:
        with zipfile.ZipFile(zip_full_path, 'r') as z:
            return len([n for n in z.namelist() if not n.endswith('/')])  # 返回文件数量
    except zipfile.BadZipFile:
        print(f"⚠️  警告: 损坏或无效的 ZIP 文件: {zip_full_path}")
        return -1  # 标记为损坏
    except Exception as e:
        print(f"⚠️  无法读取 ZIP 文件 {zip_full_path}: {e}")
        return -1

def main():
    args = get_args()

    # 检查 scan-dir 是否存在
    if not os.path.isdir(args.scan_dir):
        print(f"错误: 扫描目录不存在或不是目录: {args.scan_dir}")
        return 1

    try:
        zip_records = get_zip_info_from_db(args.db)
    except Exception as e:
        print(e)
        return 1

    print("开始检查 ZIP 文件...")

    mismatch_found = False

    for zip_path, zip_name, db_num in zip_records:
        # 构造 ZIP 文件在 --scan-dir 下的实际路径
        # zip_path 是数据库中存储的路径（如：/path/to/zip），我们只取最后一级作为目录名
        # 假设 zip_path 是相对路径或绝对路径，我们只关心最后一级目录名 + zip_name
        zip_dir_name = os.path.basename(zip_path.rstrip("/\\"))  # 提取目录名
        expected_zip_dir = os.path.join(args.scan_dir, zip_dir_name)
        zip_full_path = os.path.join(expected_zip_dir, zip_name)

        if not os.path.exists(zip_full_path):
            print(f"❌ 未找到 ZIP 文件: {zip_full_path}")
            print(f"    数据库记录数量: {db_num}")
            mismatch_found = True
            continue

        file_num = count_zip_files(zip_full_path)
        if file_num == -1:
            # 已在函数中打印错误
            mismatch_found = True
            continue

        if file_num != db_num:
            print(f"⚠️ 数量不一致: {zip_full_path}")
            print(f"    数据库记录数: {db_num}, 实际文件数: {file_num}")
            mismatch_found = True

    if not mismatch_found:
        print("✅ 所有 ZIP 文件检查通过，数量一致。")
    else:
        print("✅ 检查完成，存在不一致或缺失文件。")

    return 0 if not mismatch_found else 1

if __name__ == '__main__':
    exit(main())